# vim: expandtab:ts=4:sw=4
from __future__ import division, print_function, absolute_import

import argparse
import os

import cv2
import numpy as np
import xml.etree.ElementTree as ET
import xml.dom.minidom as minidom
import glob
import socket

from application_util import preprocessing
from application_util import visualization
from deep_track import nn_matching
from deep_track.detection import Detection
from deep_track.tracker import Tracker
import time

REPO_DIRNAME = os.path.dirname(os.path.abspath(__file__))


def gather_sequence_info(image_dir, detection_file):
    #image_dir = os.path.join(sequence_dir, "image")
    image_filenames = {
        int(os.path.splitext(f)[0]): os.path.join(image_dir, f)
        for f in os.listdir(image_dir)}

    detections = None
    if detection_file is not None:
        detections = np.load(detection_file)

    def parse_xml(filename):
        tree = ET.parse(filename)
        objs = tree.findall('object')
        num_objs = len(objs)

        boxes = []
        for ix, obj in enumerate(objs):
            bbox = obj.find('bndbox')
            x1 = float(bbox.find('xmin').text)
            y1 = float(bbox.find('ymin').text)
            x2 = float(bbox.find('xmax').text)
            y2 = float(bbox.find('ymax').text)
            boxes.append([x1, y1, x2, y2])
        return boxes

    #get door areas and car areas
    #door_xml = os.path.join(sequence_dir, "door_areas.xml")
    #door_areas = parse_xml(door_xml)
    #car_xml = os.path.join(sequence_dir, "car_areas.xml")
    #car_areas = parse_xml(car_xml)

        
    if len(image_filenames) > 0:
        image = cv2.imread(next(iter(image_filenames.values())),
                           cv2.IMREAD_GRAYSCALE)
        image_size = image.shape
    else:
        image_size = None
        
    if len(image_filenames) > 0:
        min_frame_idx = min(image_filenames.keys())
        max_frame_idx = max(image_filenames.keys())
    else:
        min_frame_idx = int(detections[:, 0].min())
        max_frame_idx = int(detections[:, 0].max())
        
    info_filename = os.path.join(image_dir, "seqinfo.ini")
    if os.path.exists(info_filename):
        with open(info_filename, "r") as f:
            line_splits = [l.split('=') for l in f.read().splitlines()[1:]]
            info_dict = dict(
                s for s in line_splits if isinstance(s, list) and len(s) == 2)

        update_ms = 1000 / int(info_dict["frameRate"])
    else:
        update_ms = None    

    feature_dim = detections.shape[1] - 10 if detections is not None else 0
    seq_info = {
        "sequence_name": os.path.basename(image_dir),
        "image_filenames": image_filenames,
        "detections": detections,
        "image_size": image_size,
        "min_frame_idx": min_frame_idx,
        "max_frame_idx": max_frame_idx,
        "feature_dim": feature_dim,
        "update_ms": update_ms
        #"door_areas": door_areas,
        #"car_areas": car_areas
    }
    return seq_info

def create_detections(detection_mat, frame_idx, min_height=0):
    frame_indices = detection_mat[:, 0].astype(np.int)
    mask = frame_indices == frame_idx

    detection_list = []
    for row in detection_mat[mask]:
        bbox, confidence, cls, feature = row[2:6], row[6], row[9], row[10:]
        #print("cls:", cls)
        if bbox[3] < min_height:
            continue
        detection_list.append(Detection(bbox, confidence, cls, feature))
    return detection_list

def get_host_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
    finally:
        s.close()

    return ip

def run(detection_dir, output_file, min_confidence,
        nms_max_overlap, min_detection_height, max_cosine_distance,
        nn_budget):
    #seq_info = gather_sequence_info(image_dir, detection_file)
    output_dir = os.path.dirname(output_file)
    if not os.path.exists(output_dir):
            os.makedirs(output_dir)
    metric = nn_matching.NearestNeighborDistanceMetric(
        "cosine", max_cosine_distance, nn_budget)
    #door_areas = seq_info["door_areas"]
    #car_areas = seq_info["car_areas"]
    result_dir = os.path.join(REPO_DIRNAME, 'result')
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
    myip = get_host_ip().split('.')[-1]
    

    tracker = Tracker(metric)
    results = []
    detection_files = glob.glob(os.path.join(detection_dir, '*.npy'))
    #print(detection_files)
    detection_files.sort()
    frame_count = 0
    for detection_file in detection_files:
    	try:
        	now_detections = np.load(detection_file)
        except Exception:
        	continue
        if len(now_detections.shape) < 2:
            for i in range(0, 40):
                tracker.predict()
            frame_count += 40
            continue
        min_frame_idx = 1
        max_frame_idx = 40
        print(detection_file)
        timecount = int(os.path.basename(detection_file).split('.')[0]) * 1000
        for frame_idx in range(min_frame_idx, max_frame_idx+1):
            detections = create_detections(
                now_detections, frame_idx, min_detection_height)
            if 0 == len(detections):
                continue
            # Run non-maxima suppression.
            boxes = np.array([d.tlwh for d in detections])
            scores = np.array([d.confidence for d in detections])
            indices = preprocessing.non_max_suppression(
                boxes, nms_max_overlap, scores)
            detections = [detections[i] for i in indices]


            # Update tracker.
            # To change
            tracker.predict()
            tracker.update(detections)

            # Store results.
            for track in tracker.tracks:
                if not track.is_confirmed() or track.time_since_update > 1:
                    continue
                bbox = track.to_tlwh()
                results.append([
                    timecount+frame_idx*250, track.track_id, bbox[0], 
                    bbox[1], bbox[2], bbox[3], int(os.path.basename(detection_file).split('.')[0]),
                    track.cls])
        frame_count += 40

    # Store results.
    f = open(output_file, 'w')
    print('%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s' % (
        'frame','tmp_id','X_min','Y_min','w',
        'h','a','b','unix_time','is_customer','true_id','box'),file=f)
    for row in results:
        print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,%d,%d,-1,-1' % (
            row[0], row[1], row[2], row[3], row[4], row[5], row[6],row[7]),file=f)

    result_file = open(os.path.join(result_dir,'track_result_'+myip+'.txt'),'a+')
    result_file.write(detection_dir)
    result_file.write('\n')
    result_file.flush()

def parse_args():
    """ Parse command line arguments.
    """
    parser = argparse.ArgumentParser(description="Deep TRACK")
    parser.add_argument(
        "--detection_dir", help="Path to custom detections.", default=None,
        required=True)
    parser.add_argument(
        "--output_file", help="Path to the tracking output file. This file will"
        " contain the tracking results on completion.",
        default="/tmp/hypotheses.txt")
    parser.add_argument(
        "--min_confidence", help="Detection confidence threshold. Disregard "
        "all detections that have a confidence lower than this value.",
        default=0.8, type=float)
    parser.add_argument(
        "--min_detection_height", help="Threshold on the detection bounding "
        "box height. Detections with height smaller than this value are "
        "disregarded", default=0, type=int)
    parser.add_argument(
        "--nms_max_overlap",  help="Non-maxima suppression threshold: Maximum "
        "detection overlap.", default=1.0, type=float)
    parser.add_argument(
        "--max_cosine_distance", help="Gating threshold for cosine distance "
        "metric (object appearance).", type=float, default=0.2)
    parser.add_argument(
        "--nn_budget", help="Maximum size of the appearance descriptors "
        "gallery. If None, no budget is enforced.", type=int, default=None)
    parser.add_argument(
        "--display", help="Show intermediate tracking results",
        default=False, type=bool)
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    run(
        args.detection_dir, args.output_file,
        args.min_confidence, args.nms_max_overlap, args.min_detection_height,
        args.max_cosine_distance, args.nn_budget)
